use std::{cmp::Ordering, collections::HashMap, fs::read_to_string};

use crate::Puzzle;

struct Point {
    x: u32,
    y: u32,
}

impl From<&str> for Point {
    fn from(s: &str) -> Self {
        let coords: Vec<&str> = s.trim().split(',').collect();
        let x = coords[0].parse().unwrap();
        let y = coords[1].parse().unwrap();
        Point { x, y }
    }
}

struct Edge(Point, Point);

pub struct HydrothermalVenture {
    edges: Vec<Edge>,
}

impl HydrothermalVenture {
    fn from(raw_data: &str) -> Self {
        let edges = raw_data
            .lines()
            .map(|line| {
                let raw_points = line.split("->").collect::<Vec<&str>>();
                let point1 = raw_points.get(0).unwrap();
                let point2 = raw_points.get(1).unwrap();
                Edge(Point::from(*point1), Point::from(*point2))
            })
            .collect::<Vec<Edge>>();

        HydrothermalVenture { edges }
    }

    fn count(&self, include_diagonal: bool) -> u32 {
        // rowIndex -> columnIndex -> count
        let mut matrix: HashMap<u32, HashMap<u32, u8>> = HashMap::new();

        // insert
        for Edge(point1, point2) in &self.edges {
            if point1.x == point2.x {
                let row_index = point1.x;
                let row = matrix.entry(row_index).or_default();
                let (start_y, end_y) = (point1.y.min(point2.y), point1.y.max(point2.y));
                for column_index in start_y..=end_y {
                    row.entry(column_index)
                        .and_modify(|count| *count += 1)
                        .or_insert(1);
                }
            } else if point1.y == point2.y {
                let column_index = point1.y;
                let (start_x, end_x) = (point1.x.min(point2.x), point1.x.max(point2.x));
                for row_index in start_x..=end_x {
                    matrix
                        .entry(row_index)
                        .or_default()
                        .entry(column_index)
                        .and_modify(|count| *count += 1)
                        .or_insert(1);
                }
            } else {
                if !include_diagonal {
                    continue;
                }

                let mut x = point1.x;
                let mut y = point1.y;
                let mutate = match (point1.x.cmp(&point2.x), point1.y.cmp(&point2.y)) {
                    (Ordering::Less, Ordering::Less) => |x: &mut u32, y: &mut u32| {
                        *x += 1;
                        *y += 1;
                    },
                    (Ordering::Less, Ordering::Greater) => |x: &mut u32, y: &mut u32| {
                        *x += 1;
                        *y -= 1;
                    },
                    (Ordering::Greater, Ordering::Less) => |x: &mut u32, y: &mut u32| {
                        *x -= 1;
                        *y += 1;
                    },
                    (Ordering::Greater, Ordering::Greater) => |x: &mut u32, y: &mut u32| {
                        *x -= 1;
                        *y -= 1;
                    },
                    _ => unreachable!(),
                };

                while x != point2.x || y != point2.y {
                    matrix
                        .entry(x)
                        .or_default()
                        .entry(y)
                        .and_modify(|count| *count += 1)
                        .or_insert(1);
                    mutate(&mut x, &mut y);
                }

                // Handle the final point
                matrix
                    .entry(x)
                    .or_default()
                    .entry(y)
                    .and_modify(|count| *count += 1)
                    .or_insert(1);
            }
        }

        matrix.iter().fold(0, |prev_result, (_, row)| {
            prev_result + row.iter().filter(|(_, count)| **count >= 2).count() as u32
        })
    }
}

impl Puzzle for HydrothermalVenture {
    type Answer = u32;

    fn new() -> Self {
        let raw_lines = read_to_string("input/day05.md").unwrap();

        Self::from(&raw_lines)
    }

    fn solve_part1(&self) -> Self::Answer {
        self.count(false)
    }

    fn solve_part2(&self) -> Self::Answer {
        self.count(true)
    }
}

#[cfg(test)]
mod test {
    use super::*;

    fn create_venture() -> HydrothermalVenture {
        let raw = "
        0,9 -> 5,9
        8,0 -> 0,8
        9,4 -> 3,4
        2,2 -> 2,1
        7,0 -> 7,4
        6,4 -> 2,0
        0,9 -> 2,9
        3,4 -> 1,4
        0,0 -> 8,8
        5,5 -> 8,2
        ";

        HydrothermalVenture::from(raw.trim())
    }

    #[test]
    fn new() {
        let venture = create_venture();

        let expected = HydrothermalVenture {
            edges: vec![
                Edge(Point { x: 0, y: 9 }, Point { x: 5, y: 9 }),
                Edge(Point { x: 8, y: 0 }, Point { x: 0, y: 8 }),
                Edge(Point { x: 9, y: 4 }, Point { x: 3, y: 4 }),
                Edge(Point { x: 2, y: 2 }, Point { x: 2, y: 1 }),
                Edge(Point { x: 7, y: 0 }, Point { x: 7, y: 4 }),
                Edge(Point { x: 6, y: 4 }, Point { x: 2, y: 0 }),
                Edge(Point { x: 0, y: 9 }, Point { x: 2, y: 9 }),
                Edge(Point { x: 3, y: 4 }, Point { x: 1, y: 4 }),
                Edge(Point { x: 0, y: 0 }, Point { x: 8, y: 8 }),
                Edge(Point { x: 5, y: 5 }, Point { x: 8, y: 2 }),
            ],
        };

        assert_eq!(venture.edges.len(), expected.edges.len());
        for (i, (actual, expected)) in venture.edges.iter().zip(expected.edges.iter()).enumerate() {
            assert_eq!(actual.0.x, expected.0.x, "Edge {i} start x mismatch");
            assert_eq!(actual.0.y, expected.0.y, "Edge {i} start y mismatch");
            assert_eq!(actual.1.x, expected.1.x, "Edge {i} end x mismatch");
            assert_eq!(actual.1.y, expected.1.y, "Edge {i} end y mismatch");
        }
    }

    #[test]
    fn part1() {
        assert_eq!(create_venture().solve_part1(), 5);
    }

    #[test]
    fn part2() {
        assert_eq!(create_venture().solve_part2(), 12);
    }
}
